from utils.vocabulary import get_big_vocab
from cio.jsonl import JsonLReader, JsonLWriter
from config.Config import BASE_DIR

vocab = get_big_vocab()

data = JsonLReader(BASE_DIR / 'resource' / 'data' / 'consum' / 'split_context.train.jsonl').yield_read()
writer = JsonLWriter(BASE_DIR / 'resource' / 'data' / 'consum' / 'context_words_id.train.jsonl')
for each in data:
    new_data = {
        "method": vocab.lookup_indices(each['method']),
        "summary": vocab.lookup_indices(each['summary']),
        "context": vocab.lookup_indices(each['context']),
    }
    writer.save_one(new_data)

data = JsonLReader(BASE_DIR / 'resource' / 'data' / 'consum' / 'split_context.test.jsonl').yield_read()
writer = JsonLWriter(BASE_DIR / 'resource' / 'data' / 'consum' / 'context_words_id.test.jsonl')
for each in data:
    new_data = {
        "method": vocab.lookup_indices(each['method']),
        "summary": vocab.lookup_indices(each['summary']),
        "context": vocab.lookup_indices(each['context']),
    }
    writer.save_one(new_data)
